#include <catch2/catch_test_macros.hpp>
#include "mpiHelpers.hpp"

using namespace hms;

TEST_CASE( "MPI_Gather_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	SECTION( "Single value" ){
		ArrayXi arr = MPI_Gather_H<ArrayXi,Index>(comm.rank(), comm);

		if ( comm.isMaster() ){
		for ( int i{0}; i<comm.size(); ++i )
			REQUIRE( arr(i) == i );
		}
	}
	SECTION( "Multiple values" ){
		int length {3};
		ArrayXi arr (length);
		arr = comm.rank();

		ArrayXi gathered { MPI_Gather_H( arr, comm, comm.master() ) };

		if ( comm.isMaster() ){
			for ( int i{0}; i<comm.size(); ++i )
				for ( int j{0}; j<length; ++j )
					REQUIRE( gathered( i*length+j ) == i );
		}
	}
}

TEST_CASE( "MPI_Allgather_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	SECTION( "Single value" ){
		ArrayXi arr = MPI_Allgather_H<ArrayXi,Index>(comm.rank(), comm);

		for ( int i{0}; i<comm.size(); ++i )
			REQUIRE( arr(i) == i );
	}
	SECTION( "Multiple values" ){
		int length {3};
		ArrayXi arr (length);
		arr = comm.rank();

		ArrayXi gathered { MPI_Allgather_H( arr, comm ) };

		for ( int i{0}; i<comm.size(); ++i )
			for ( int j{0}; j<length; ++j )
				REQUIRE( gathered( i*length+j ) == i );
	}
}

TEST_CASE( "MPI_Gatherv_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	/* different lengths on all ranks */
	ArrayXi arr ( comm.rank() + 1 );
	arr = comm.rank() + 1;

	ArrayXi gathered { MPI_Gatherv_H( arr, comm, comm.master() ) };

	if ( comm.isMaster() ){
		int idx {0};
		for ( int i{0}; i<comm.size(); ++i )
			for ( int j{0}; j<i+1; ++j )
				REQUIRE( gathered(idx++) == i+1 );
	}
}

TEST_CASE( "MPI_Allgatherv_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	/* different lengths on all ranks */
	ArrayXi arr ( comm.rank() + 1 );
	arr = comm.rank() + 1;

	ArrayXi gathered { MPI_Allgatherv_H( arr, comm ) };

	int idx {0};
	for ( int i{0}; i<comm.size(); ++i )
		for ( int j{0}; j<i+1; ++j )
			REQUIRE( gathered(idx++) == i+1 );
	
	SECTION( "String" ){
		std::string myName { [&](){
			std::stringstream ss;
			ss << comm.rank() << " ";
			return ss.str();
		}() };
		std::string allNames { [&](){
			std::stringstream ss;
			for ( int i{0}; i<comm.size(); ++i )
				ss << i << " ";
			return ss.str();
		}() };

		std::string allNamesTest { MPI_Allgatherv_H( myName, comm ) };
	}
}

TEST_CASE( "MPI_Bcast_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	SECTION( "Single value" ){
		char val;
		if ( comm.isMaster() ) val = 'c';
		MPI_Bcast_H( val, comm );

		REQUIRE( val == 'c' );
	}
	SECTION( "Multiple values, no resizing" ){
		Array2i val;
		if ( comm.isMaster() ) val = 123;
		MPI_Bcast_H( val, comm, comm.master(), false );

		REQUIRE( (val == 123).all() );
	}
	SECTION( "Multiple values, with resizing" ){
		ArrayXi val;
		if ( comm.isMaster() ){
			val.resize(3);
			val = 123;
		}
		MPI_Bcast_H( val, comm, comm.master(), true );

		REQUIRE( (val == 123).all() );
	}
}

TEST_CASE( "MPI_Win_H" ){
	MPI_Comm_H comm { MPI_COMM_WORLD };

	ArrayXi arr;

	if ( comm.isMaster() )
		arr.resize( comm.size() );

	/* create a window for each section */
	MPI_Win_H win ( arr, comm );

	/* each rank has an array with length 1, into which it gets or from which
	 * it puts data from/to master */
	ArrayXi val (1);

	SECTION( "get" ){
		if ( comm.isMaster() )
			arr = ArrayXi::LinSpaced( comm.size(), 0, comm.size()-1 );
		
		win.fence();
		win.get( val, 0, 1, comm.master(), comm.rank() );
		win.fence();
		
		REQUIRE( val(0) == comm.rank() );
	}
	SECTION( "put" ){
		val = comm.rank();

		win.fence();
		win.put( val, 0, 1, comm.master(), comm.rank() );
		win.fence();

		if ( comm.isMaster() )
			for ( int i{0}; i<comm.size(); ++i )
				REQUIRE( arr(i) == i );
	}
	SECTION( "override window (call create on already existing window)" ){
		ArrayXi arr2;

		if ( comm.isMaster() )
			arr2 = ArrayXi::LinSpaced( comm.size(), 0, comm.size()-1 );
		
		win.create( arr2, comm );

		win.fence();
		win.get( val, 0, 1, comm.master(), comm.rank() );
		win.fence();

		REQUIRE( val(0) == comm.rank() );
	}
}